"""
智能体协调器
协调多个智能体协作，分配任务并管理通信
"""

from typing import Dict, List, Any, Optional, Set
import asyncio
import json
from datetime import datetime
from enum import Enum

from apps.agents.base_agent import BaseAgent, AgentMessage, AgentType
from apps.agents.text_qa_agent import TextQAAgent
from apps.agents.ocr_agent import OCRAgent
from apps.agents.calculator_agent import CalculatorAgent
from utils.logger import logger
from states.agent_state import AgentState


class TaskType(Enum):
    """任务类型"""
    TEXT_QA = "text_qa"
    OCR = "ocr"
    CALCULATION = "calculation"
    MULTI_MODAL = "multi_modal"
    WORKFLOW = "workflow"


class OrchestratorDecision(Enum):
    """协调器决策类型"""
    SINGLE_AGENT = "single_agent"
    MULTI_AGENT_PARALLEL = "multi_agent_parallel"
    MULTI_AGENT_SEQUENTIAL = "multi_agent_sequential"
    WORKFLOW = "workflow"
    FALLBACK = "fallback"


class AgentOrchestrator:
    """智能体协调器"""
    
    def __init__(self):
        self.agents: Dict[str, BaseAgent] = {}
        self.agent_states: Dict[str, AgentState] = {}
        self.agent_capabilities: Dict[AgentType, List[str]] = {}
        self.task_queue: asyncio.Queue = asyncio.Queue()
        self.message_bus: Dict[str, List[AgentMessage]] = {}
        self.active_tasks: Dict[str, Dict[str, Any]] = {}
        self.task_results: Dict[str, Any] = {}
        self.running = True
        
        # 初始化智能体
        self._initialize_agents()
        
        # 定义智能体能力
        self._define_agent_capabilities()
        
        # 启动心跳维护
        self._start_heartbeat_maintenance()
        
        logger.info("智能体协调器初始化完成")
    
    def _initialize_agents(self):
        """初始化所有智能体并注册到Redis"""
        
        # 文本问答智能体
        text_qa_agent = TextQAAgent("text_qa_agent_1")
        self.agents[text_qa_agent.agent_id] = text_qa_agent
        
        # 创建Redis客户端并验证连接
        from middleware.redis_adapter import get_redis_adapter
        from auxiliary_agents.health_monitor import HealthCheckAgent
        redis_client = get_redis_adapter()
        
        # 获取HealthCheckAgent使用的前缀
        agent_prefix = getattr(HealthCheckAgent, 'agent_prefix', 'agent:')
        logger.info(f"HealthCheckAgent前缀: {agent_prefix}")
        
        # 验证Redis连接
        try:
            redis_ping = redis_client.ping()
            logger.info(f"Redis连接状态: {redis_ping}")
        except Exception as e:
            logger.error(f"Redis连接失败: {str(e)}")
        
        # 文本问答智能体 - 使用正确的格式
        text_qa_state = AgentState(
            agent_type=AgentType.TEXT_QA.value,
            instance_id=text_qa_agent.agent_id,
            port=8001
        )
        self.agent_states[text_qa_agent.agent_id] = text_qa_state
        
        # 手动添加到Redis集合中并设置正确的状态
        try:
            # 根据health_monitor的测试代码，添加_agent后缀
            agent_type_with_suffix = f"{AgentType.TEXT_QA.value}_agent"
            agent_ids_key = f"agent:ids:{agent_type_with_suffix}"
            agent_state_key = f"{agent_prefix}{text_qa_agent.agent_id}"
            
            logger.info(f"正在注册智能体: agent_type={agent_type_with_suffix}, agent_id={text_qa_agent.agent_id}")
            logger.info(f"使用的键: agent_ids_key={agent_ids_key}, agent_state_key={agent_state_key}")
            
            # 添加到集合
            sadd_result = redis_client.sadd(agent_ids_key, text_qa_agent.agent_id)
            logger.info(f"添加到Redis集合结果: {sadd_result}，键: {agent_ids_key}")
            
            # 设置状态
            hset_result = redis_client.hset(agent_state_key, "status", "UP")
            logger.info(f"设置状态结果: {hset_result}，键: {agent_state_key}")
            
            redis_client.hset(agent_state_key, "task_count", "0")
            redis_client.expire(agent_state_key, 1800)
            
            # 验证设置是否成功
            status = redis_client.hget(agent_state_key, "status")
            in_set = redis_client.sismember(agent_ids_key, text_qa_agent.agent_id)
            logger.info(f"TextQAAgent Redis验证 - status={status}, in_set={in_set}")
        except Exception as e:
            logger.error(f"TextQAAgent Redis操作失败: {str(e)}", exc_info=True)
        
        # OCR智能体
        ocr_agent = OCRAgent("ocr_agent_1")
        self.agents[ocr_agent.agent_id] = ocr_agent
        
        ocr_state = AgentState(
            agent_type=AgentType.MULTI_MODAL.value,
            instance_id=ocr_agent.agent_id,
            port=8002
        )
        self.agent_states[ocr_agent.agent_id] = ocr_state
        
        try:
            agent_type_with_suffix = f"{AgentType.MULTI_MODAL.value}_agent"
            agent_ids_key = f"agent:ids:{agent_type_with_suffix}"
            agent_state_key = f"{agent_prefix}{ocr_agent.agent_id}"
            
            logger.info(f"正在注册智能体: agent_type={agent_type_with_suffix}, agent_id={ocr_agent.agent_id}")
            
            redis_client.sadd(agent_ids_key, ocr_agent.agent_id)
            redis_client.hset(agent_state_key, "status", "UP")
            redis_client.hset(agent_state_key, "task_count", "0")
            redis_client.expire(agent_state_key, 1800)
            logger.info(f"OCRAgent Redis操作成功")
        except Exception as e:
            logger.error(f"OCRAgent Redis操作失败: {str(e)}", exc_info=True)
        
        # 计算器智能体
        calculator_agent = CalculatorAgent("calculator_agent_1")
        self.agents[calculator_agent.agent_id] = calculator_agent
        
        calculator_state = AgentState(
            agent_type=AgentType.TOOL_CALL.value,
            instance_id=calculator_agent.agent_id,
            port=8003
        )
        self.agent_states[calculator_agent.agent_id] = calculator_state
        
        try:
            agent_type_with_suffix = f"{AgentType.TOOL_CALL.value}_agent"
            agent_ids_key = f"agent:ids:{agent_type_with_suffix}"
            agent_state_key = f"{agent_prefix}{calculator_agent.agent_id}"
            
            logger.info(f"正在注册智能体: agent_type={agent_type_with_suffix}, agent_id={calculator_agent.agent_id}")
            
            redis_client.sadd(agent_ids_key, calculator_agent.agent_id)
            redis_client.hset(agent_state_key, "status", "UP")
            redis_client.hset(agent_state_key, "task_count", "0")
            redis_client.expire(agent_state_key, 1800)
            logger.info(f"CalculatorAgent Redis操作成功")
        except Exception as e:
            logger.error(f"CalculatorAgent Redis操作失败: {str(e)}", exc_info=True)
        
        logger.info(f"已初始化并注册 {len(self.agents)} 个智能体")
    
    def _define_agent_capabilities(self):
        """定义智能体能力"""
        
        self.agent_capabilities = {
            AgentType.TEXT_QA: [
                "text_analysis", "question_answering", "summarization",
                "fact_checking", "comparison", "text_classification"
            ],
            AgentType.OCR: [
                "image_text_extraction", "document_scanning",
                "handwriting_recognition", "batch_processing"
            ],
            AgentType.CALCULATOR: [
                "basic_calculations", "scientific_calculations",
                "statistical_analysis", "financial_calculations",
                "unit_conversions", "expression_evaluation"
            ]
        }
    
    async def analyze_task(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """分析任务需求"""
        
        logger.info("开始分析任务需求")
        
        # 提取任务特征
        task_features = self._extract_task_features(task_data)
        
        # 匹配智能体能力
        matched_agents = self._match_agent_capabilities(task_features)
        
        # 生成决策建议
        decision = self._generate_decision(task_features, matched_agents)
        
        analysis_result = {
            "task_features": task_features,
            "matched_agents": matched_agents,
            "decision": decision,
            "confidence": self._calculate_confidence(task_features, matched_agents),
            "timestamp": datetime.now().isoformat()
        }
        
        logger.info(f"任务分析完成: {decision['type']}")
        return analysis_result
    
    def _extract_task_features(self, task_data: Dict[str, Any]) -> Dict[str, Any]:
        """提取任务特征"""
        
        features = {
            "input_types": [],
            "output_requirements": [],
            "complexity_level": "medium",
            "domain": "general",
            "constraints": []
        }
        
        # 分析输入类型
        if "text" in task_data:
            features["input_types"].append("text")
            
            # 分析文本特征
            text = task_data["text"]
            if len(text) > 1000:
                features["complexity_level"] = "high"
            
            # 检查是否包含问题
            if any(q in text.lower() for q in ["?", "什么", "如何", "为什么", "哪里"]):
                features["domain"] = "qa"
        
        if "image" in task_data or "image_path" in task_data:
            features["input_types"].append("image")
            
            # 检查图像处理需求
            if "extract_text" in task_data.get("requirements", []):
                features["domain"] = "ocr"
        
        if "calculation" in task_data or "expression" in task_data:
            features["input_types"].append("numerical")
            features["domain"] = "calculation"
        
        # 分析输出要求
        if "requirements" in task_data:
            features["output_requirements"] = task_data["requirements"]
        
        # 分析约束条件
        if "constraints" in task_data:
            features["constraints"] = task_data["constraints"]
        
        return features
    
    def _match_agent_capabilities(self, task_features: Dict[str, Any]) -> List[Dict[str, Any]]:
        """匹配智能体能力"""
        
        matched_agents = []
        
        for agent_id, agent in self.agents.items():
            agent_capabilities = self.agent_capabilities.get(agent.agent_type, [])
            
            # 计算匹配分数
            match_score = self._calculate_match_score(
                task_features, agent_capabilities, agent
            )
            
            if match_score > 0:
                matched_agents.append({
                    "agent_id": agent_id,
                    "agent_type": agent.agent_type,
                    "match_score": match_score,
                    "capabilities": agent_capabilities
                })
        
        # 按匹配分数排序
        matched_agents.sort(key=lambda x: x["match_score"], reverse=True)
        
        return matched_agents
    
    def _start_heartbeat_maintenance(self):
        """启动智能体心跳维护"""
        logger.info("智能体心跳维护机制已启动")
        
        # 在异步环境中，我们不再使用单独的线程，而是依赖AgentState内部的心跳机制
        # 每个AgentState实例在初始化时会自动处理心跳注册和更新
        for agent_id, agent_state in self.agent_states.items():
            if hasattr(agent_state, 'send_heartbeat'):
                # 立即发送一次心跳以确保注册成功
                try:
                    agent_state.send_heartbeat()
                    logger.info(f"立即发送智能体心跳: {agent_id}")
                except Exception as e:
                    logger.error(f"发送智能体心跳失败: {agent_id}, 错误: {str(e)}")
    
    def shutdown(self):
        """关闭所有智能体并注销Redis注册"""
        try:
            self.running = False
            
            # 先注销Redis注册
            for agent_id, agent_state in self.agent_states.items():
                if hasattr(agent_state, "unregister"):
                    try:
                        agent_state.unregister()
                        logger.info(f"注销智能体Redis注册: {agent_id}")
                    except Exception as e:
                        logger.error(f"注销智能体Redis注册失败: {agent_id}, 错误: {str(e)}")
            
            # 再关闭智能体
            for agent_id, agent in self.agents.items():
                if hasattr(agent, "shutdown"):
                    try:
                        agent.shutdown()
                        logger.info(f"关闭智能体: {agent_id}")
                    except Exception as e:
                        logger.error(f"关闭智能体失败: {agent_id}, 错误: {str(e)}")
            
            logger.info("所有智能体已关闭")
        except Exception as e:
            logger.error(f"关闭智能体协调器失败: {str(e)}")
    
    def _calculate_match_score(self, task_features: Dict[str, Any], 
                              agent_capabilities: List[str], agent: BaseAgent) -> float:
        """计算匹配分数"""
        
        score = 0.0
        
        # 基于输入类型的匹配
        input_types = task_features.get("input_types", [])
        
        if "text" in input_types and agent.agent_type == AgentType.TEXT_QA:
            score += 0.4
        elif "image" in input_types and agent.agent_type == AgentType.OCR:
            score += 0.4
        elif "numerical" in input_types and agent.agent_type == AgentType.CALCULATOR:
            score += 0.4
        
        # 基于输出要求的匹配
        output_requirements = task_features.get("output_requirements", [])
        
        for requirement in output_requirements:
            if requirement in agent_capabilities:
                score += 0.2
        
        # 基于域的匹配
        domain = task_features.get("domain", "general")
        
        if domain == "qa" and agent.agent_type == AgentType.TEXT_QA:
            score += 0.2
        elif domain == "ocr" and agent.agent_type == AgentType.OCR:
            score += 0.2
        elif domain == "calculation" and agent.agent_type == AgentType.CALCULATOR:
            score += 0.2
        
        return min(score, 1.0)  # 确保分数不超过1.0
    
    def _generate_decision(self, task_features: Dict[str, Any], 
                          matched_agents: List[Dict[str, Any]]) -> Dict[str, Any]:
        """生成决策"""
        
        if not matched_agents:
            return {
                "type": OrchestratorDecision.FALLBACK,
                "reasoning": "没有匹配的智能体",
                "recommended_agents": []
            }
        
        # 单智能体决策
        if len(matched_agents) == 1:
            return {
                "type": OrchestratorDecision.SINGLE_AGENT,
                "reasoning": "单个智能体完全匹配任务需求",
                "recommended_agents": [matched_agents[0]["agent_id"]]
            }
        
        # 多智能体并行决策
        if len(matched_agents) > 1 and self._can_parallel_execute(task_features):
            return {
                "type": OrchestratorDecision.MULTI_AGENT_PARALLEL,
                "reasoning": "多个智能体可以并行执行",
                "recommended_agents": [agent["agent_id"] for agent in matched_agents]
            }
        
        # 多智能体串行决策
        if len(matched_agents) > 1:
            return {
                "type": OrchestratorDecision.MULTI_AGENT_SEQUENTIAL,
                "reasoning": "多个智能体需要串行执行",
                "recommended_agents": [agent["agent_id"] for agent in matched_agents]
            }
        
        return {
            "type": OrchestratorDecision.FALLBACK,
            "reasoning": "无法确定最佳执行策略",
            "recommended_agents": []
        }
    
    def _can_parallel_execute(self, task_features: Dict[str, Any]) -> bool:
        """判断是否可以并行执行"""
        
        # 检查约束条件
        constraints = task_features.get("constraints", [])
        
        if "sequential" in constraints:
            return False
        
        # 检查复杂度
        complexity = task_features.get("complexity_level", "medium")
        
        if complexity == "high":
            return False
        
        return True
    
    def _calculate_confidence(self, task_features: Dict[str, Any], 
                             matched_agents: List[Dict[str, Any]]) -> float:
        """计算决策置信度"""
        
        if not matched_agents:
            return 0.0
        
        # 基于匹配分数的置信度
        max_match_score = max(agent["match_score"] for agent in matched_agents)
        
        # 基于智能体数量的置信度
        agent_count_factor = min(len(matched_agents) * 0.1, 0.3)
        
        confidence = max_match_score + agent_count_factor
        
        return min(confidence, 1.0)
    
    async def execute_task(self, task_data: Dict[str, Any], 
                          decision: Dict[str, Any]) -> Dict[str, Any]:
        """执行任务"""
        
        task_id = f"task_{asyncio.get_event_loop().time()}"
        
        logger.info(f"开始执行任务: {task_id}, 决策类型: {decision['type']}")
        
        try:
            # 记录活跃任务
            self.active_tasks[task_id] = {
                "task_data": task_data,
                "decision": decision,
                "start_time": datetime.now(),
                "status": "running"
            }
            
            # 根据决策类型执行
            if decision["type"] == OrchestratorDecision.SINGLE_AGENT:
                result = await self._execute_single_agent(task_id, task_data, decision)
                
            elif decision["type"] == OrchestratorDecision.MULTI_AGENT_PARALLEL:
                result = await self._execute_multi_agent_parallel(task_id, task_data, decision)
                
            elif decision["type"] == OrchestratorDecision.MULTI_AGENT_SEQUENTIAL:
                result = await self._execute_multi_agent_sequential(task_id, task_data, decision)
                
            else:
                result = await self._execute_fallback(task_id, task_data, decision)
            
            # 记录结果
            self.task_results[task_id] = result
            self.active_tasks[task_id]["status"] = "completed"
            self.active_tasks[task_id]["end_time"] = datetime.now()
            
            logger.info(f"任务执行完成: {task_id}")
            return result
            
        except Exception as e:
            logger.error(f"任务执行失败: {task_id}, 错误: {str(e)}")
            self.active_tasks[task_id]["status"] = "failed"
            self.active_tasks[task_id]["error"] = str(e)
            
            raise
    
    async def _execute_single_agent(self, task_id: str, task_data: Dict[str, Any], 
                                   decision: Dict[str, Any]) -> Dict[str, Any]:
        """执行单智能体任务"""
        
        agent_id = decision["recommended_agents"][0]
        agent = self.agents[agent_id]
        
        logger.info(f"执行单智能体任务: {agent_id}")
        
        # 发送任务给智能体
        message = AgentMessage(
            sender="orchestrator",
            receiver=agent_id,
            message_type="task_assignment",
            content=task_data
        )
        
        # 处理任务
        result = await agent.process_task(task_data)
        
        return {
            "task_id": task_id,
            "execution_type": "single_agent",
            "agent_id": agent_id,
            "result": result,
            "timestamp": datetime.now().isoformat()
        }
    
    async def _execute_multi_agent_parallel(self, task_id: str, task_data: Dict[str, Any], 
                                           decision: Dict[str, Any]) -> Dict[str, Any]:
        """执行多智能体并行任务"""
        
        agent_ids = decision["recommended_agents"]
        
        logger.info(f"执行多智能体并行任务: {agent_ids}")
        
        # 创建并行任务
        tasks = []
        for agent_id in agent_ids:
            agent = self.agents[agent_id]
            task = agent.process_task(task_data)
            tasks.append(task)
        
        # 并行执行
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # 处理结果
        processed_results = []
        for i, result in enumerate(results):
            agent_id = agent_ids[i]
            
            if isinstance(result, Exception):
                processed_results.append({
                    "agent_id": agent_id,
                    "status": "failed",
                    "error": str(result)
                })
            else:
                processed_results.append({
                    "agent_id": agent_id,
                    "status": "success",
                    "result": result
                })
        
        return {
            "task_id": task_id,
            "execution_type": "multi_agent_parallel",
            "agent_results": processed_results,
            "timestamp": datetime.now().isoformat()
        }
    
    async def _execute_multi_agent_sequential(self, task_id: str, task_data: Dict[str, Any], 
                                             decision: Dict[str, Any]) -> Dict[str, Any]:
        """执行多智能体串行任务"""
        
        agent_ids = decision["recommended_agents"]
        
        logger.info(f"执行多智能体串行任务: {agent_ids}")
        
        results = []
        current_data = task_data.copy()
        
        for agent_id in agent_ids:
            agent = self.agents[agent_id]
            
            try:
                # 执行当前智能体
                result = await agent.process_task(current_data)
                
                results.append({
                    "agent_id": agent_id,
                    "status": "success",
                    "result": result
                })
                
                # 更新数据供下一个智能体使用
                current_data["previous_result"] = result
                
            except Exception as e:
                results.append({
                    "agent_id": agent_id,
                    "status": "failed",
                    "error": str(e)
                })
                
                # 如果某个智能体失败，停止后续执行
                break
        
        return {
            "task_id": task_id,
            "execution_type": "multi_agent_sequential",
            "agent_results": results,
            "timestamp": datetime.now().isoformat()
        }
    
    async def _execute_fallback(self, task_id: str, task_data: Dict[str, Any], 
                               decision: Dict[str, Any]) -> Dict[str, Any]:
        """执行降级方案"""
        
        logger.warning(f"执行降级方案: {task_id}")
        
        # 尝试使用默认智能体
        default_agent_id = list(self.agents.keys())[0]
        default_agent = self.agents[default_agent_id]
        
        try:
            result = await default_agent.process_task(task_data)
            
            return {
                "task_id": task_id,
                "execution_type": "fallback",
                "agent_id": default_agent_id,
                "result": result,
                "warning": "使用降级方案执行",
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            logger.error(f"降级方案也失败了: {str(e)}")
            
            return {
                "task_id": task_id,
                "execution_type": "fallback_failed",
                "error": str(e),
                "timestamp": datetime.now().isoformat()
            }
    
    async def send_message(self, message: AgentMessage):
        """发送消息"""
        
        # 添加到消息总线
        if message.receiver not in self.message_bus:
            self.message_bus[message.receiver] = []
        
        self.message_bus[message.receiver].append(message)
        
        logger.info(f"消息已发送: {message.message_id} -> {message.receiver}")
    
    async def get_messages(self, agent_id: str) -> List[AgentMessage]:
        """获取智能体的消息"""
        
        messages = self.message_bus.get(agent_id, [])
        
        # 清空已获取的消息
        if agent_id in self.message_bus:
            del self.message_bus[agent_id]
        
        return messages
    
    async def get_agent_status(self, agent_id: str) -> Dict[str, Any]:
        """获取智能体状态"""
        
        agent = self.agents.get(agent_id)
        if not agent:
            return {"status": "not_found"}
        
        return {
            "agent_id": agent_id,
            "agent_type": agent.agent_type,
            "status": agent.status,
            "health_status": agent.health_status,
            "last_activity": agent.last_activity
        }
    
    async def get_system_status(self) -> Dict[str, Any]:
        """获取系统状态"""
        
        agent_statuses = {}
        for agent_id in self.agents:
            agent_statuses[agent_id] = await self.get_agent_status(agent_id)
        
        return {
            "total_agents": len(self.agents),
            "active_tasks": len(self.active_tasks),
            "completed_tasks": len(self.task_results),
            "agent_statuses": agent_statuses,
            "timestamp": datetime.now().isoformat()
        }
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        
        healthy_agents = 0
        total_agents = len(self.agents)
        
        for agent_id, agent in self.agents.items():
            agent_health = await agent.health_check()
            if agent_health.get("status") == "healthy":
                healthy_agents += 1
        
        return {
            "orchestrator_status": "healthy",
            "total_agents": total_agents,
            "healthy_agents": healthy_agents,
            "health_percentage": (healthy_agents / total_agents) * 100 if total_agents > 0 else 0,
            "timestamp": datetime.now().isoformat()
        }